'''
Author: Morphlng
Date: 2022-05-05 18:20:32
LastEditTime: 2022-05-06 10:47:47
LastEditors: Morphlng
Description: 
FilePath: \flappy_bird\flappyBird_agent.py
'''

import parl
import paddle
import paddle.fluid as fluid
import numpy as np


class FlappyBirdAgent(parl.Agent):
    """Agent of FlappyBird-V0 env.
    Args:
        algorithm(parl.Algorithm): algorithm used to solve the problem.
    """

    def __init__(self, algorithm, obs_dim, act_dim, e_greed=0.1, e_greed_decrement=0):
        self.obs_dim = obs_dim
        self.act_dim = act_dim
        super(FlappyBirdAgent, self).__init__(algorithm)

        self.global_step = 0
        # copy parameters from model to target in every 200 step
        self.update_target_steps = 200
        # probability to exploring actions randomly
        self.e_greed = e_greed
        # as training gradually converges, less exploration can be implemented
        self.e_greed_decrement = e_greed_decrement

    def sample(self, obs):
        """Sample an action when given an observation
        Args:
            obs(np.float32): shape of (obs_dim,)

        Returns:
            act(int): action
        """
        obs = paddle.to_tensor(obs, dtype='float32')
        sample = np.random.rand()  # generate random number between 0-1
        if sample < self.e_greed:
            # exploration: every action may be selected
            act = np.random.randint(self.act_dim)
        else:
            act = self.predict(obs)  # choose the best action
        # as training gradually converges, less exploration can be implemented
        self.e_greed = max(0.15, self.e_greed - self.e_greed_decrement)
        return act

    def predict(self, obs):
        """Predict an action when given an observation
        Args:
            obs(np.float32): shape of (obs_dim,)

        Returns:
            act(int): action
        """
        obs = paddle.to_tensor(obs, dtype='float32')
        pred_Q = self.alg.predict(obs)

        # choose the index of max Q value: the corresponding action
        act = pred_Q.argmax().numpy()[0]
        return act

    def learn(self, obs, act, reward, next_obs, terminal):
        """Update model with an episode data
        Args:
            obs(np.float32): shape of (batch_size, obs_dim)
            act(np.int32): shape of (batch_size)
            reward(np.float32): shape of (batch_size)
            next_obs(np.float32): shape of (batch_size, obs_dim)
            terminal(np.bool): shape of (batch_size)

        Returns:
            loss(float)
        """
        # copy parameters from model to target model in every 200 training steps
        if self.global_step % self.update_target_steps == 0:
            self.alg.sync_target()
        self.global_step += 1

        act = np.expand_dims(act, -1)
        reward = np.expand_dims(reward, -1)
        terminal = np.expand_dims(terminal, -1)

        obs = paddle.to_tensor(obs, dtype='float32')
        act = paddle.to_tensor(act, dtype='int32')
        reward = paddle.to_tensor(reward, dtype='float32')
        next_obs = paddle.to_tensor(next_obs, dtype='float32')
        terminal = paddle.to_tensor(terminal, dtype='int32')

        loss = self.alg.learn(obs, act, reward, next_obs, terminal)
        return loss
